# -*- coding: utf-8 -*-
"""
PyQt6 全新界面重构
===================
左侧: 图片列表 + 浏览按钮
右侧: 参数设置、背景设置、算法选择
底部: 自动裁剪、开始处理按钮、进度条

* 完全保留原有所有功能、按钮、选项, 仅改变布局与界面库。
* 所有算法调用(distance_alpha / floodfill_alpha)未改动。
* 提供 SimpleVar / SliderWrapper 等兼容 tkinter 调用, 保证后续业务代码无需改。
* RegionSelectionDialog 使用 QRubberBand 框选颜色; sample_bg_color 支持平均取色。
"""

import sys
import os
import threading
from collections import Counter
from typing import Callable, List, Tuple

import cv2
import numpy as np
from PIL import Image, ImageQt
from PyQt6 import QtCore, QtGui, QtWidgets

from detector_color import distance_alpha  # 业务算法, 保持不变
from detector_portrait import floodfill_alpha  # 业务算法, 保持不变
from detector_ai import u2net_alpha  # AI 抠图算法 (U²-Net)

# --------------------- 兼容包装 ---------------------
class SimpleVar:
    def __init__(self, value=None):
        self._value = value

    def get(self):
        return self._value

    def set(self, value):
        self._value = value


class SliderWrapper:
    def __init__(self, slider: QtWidgets.QSlider):
        self.slider = slider

    def get(self):
        return self.slider.value()

    def set(self, val):
        self.slider.setValue(int(val))


class CheckBoxWrapper:
    def __init__(self, checkbox: QtWidgets.QCheckBox):
        self.checkbox = checkbox

    def get(self):
        return self.checkbox.isChecked()

    def set(self, val):
        self.checkbox.setChecked(bool(val))


class RootAdapter(QtCore.QObject):
    """提供 after/update 接口, 兼容原 tkinter 调用 (线程安全)"""

    @QtCore.pyqtSlot(object)
    def _exec_in_gui(self, func: Callable):
        """在主线程执行传入函数"""
        try:
            func()
        except Exception as err:  # noqa: BLE001
            print("GUI callback error:", err)

    def after(self, delay_ms: int, func: Callable):
        if delay_ms == 0:
            # 使用 invokeMethod 保证在 GUI 线程执行
            QtCore.QMetaObject.invokeMethod(
                self,
                "_exec_in_gui",
                QtCore.Qt.ConnectionType.QueuedConnection,
                QtCore.Q_ARG(object, func),
            )
        else:
            QtCore.QTimer.singleShot(delay_ms, lambda: self.after(0, func))

    def update(self):
        QtWidgets.QApplication.processEvents()


# --------------------- 取样对话框 ---------------------
class RegionSelectionDialog(QtWidgets.QDialog):
    def __init__(self, image_path: str, callback: Callable, *, average: bool, parent=None):
        super().__init__(parent)
        self.setWindowTitle("拖拽框选区域")
        self.callback = callback
        self._average = average
        self._origin = QtCore.QPoint()
        self._band: QtWidgets.QRubberBand | None = None
        self._confirmed_rects: list[QtWidgets.QWidget] = []  # 保存已确认矩形
        self._img_lbl: QtWidgets.QLabel | None = None

        # 加载图像并缩放 (最长 800)
        try:
            # 直接用 QPixmap 可以避免部分编码问题
            pix = QtGui.QPixmap(image_path)
            if pix.isNull():
                # 使用 PIL 退路
                pil_img = Image.open(image_path).convert("RGB")
                self._pil_img = pil_img
                qt_img = ImageQt.ImageQt(pil_img)
                pix = QtGui.QPixmap.fromImage(qt_img)
            else:
                # 需同步保存 PIL 版本
                pil_img = Image.open(image_path).convert("RGB")
                self._pil_img = pil_img
            # 计算缩放比例并按需缩放到最长边 800
            max_len = 800
            self._ratio = 1.0
            if max(pix.width(), pix.height()) > max_len:
                self._ratio = max_len / max(pix.width(), pix.height())
                pix = pix.scaled(int(pix.width()*self._ratio), int(pix.height()*self._ratio),
                                 QtCore.Qt.AspectRatioMode.KeepAspectRatio,
                                 QtCore.Qt.TransformationMode.SmoothTransformation)
        except Exception as e:
            QtWidgets.QMessageBox.critical(self, "错误", f"无法加载图片:\n{e}")
            self.reject()
            return

        lbl = QtWidgets.QLabel()
        lbl.setPixmap(pix)
        lbl.setFixedSize(pix.size())
        lbl.setMouseTracking(True)
        lbl.installEventFilter(self)

        lay = QtWidgets.QVBoxLayout(self)
        lay.addWidget(lbl)
        self._img_lbl = lbl

        # -------- 取样方式选择 (仅在非平均模式) --------
        if not average:
            opt_lay = QtWidgets.QHBoxLayout()
            lay.addLayout(opt_lay)
            opt_lay.addWidget(QtWidgets.QLabel("取样方式:"))
            self._mode_combo = QtWidgets.QComboBox()
            self._mode_combo.addItems(["高频颜色", "全部颜色"])
            self._mode_combo.setCurrentIndex(0)   # 默认选"高频前 N 色"
            opt_lay.addWidget(self._mode_combo)

            self._spin_top_n = QtWidgets.QSpinBox()
            self._spin_top_n.setRange(5, 500)
            self._spin_top_n.setValue(30)
            self._spin_top_n.setSingleStep(5)
            self._spin_top_n.setFixedWidth(70)
            opt_lay.addWidget(self._spin_top_n)

            # 只有选择"高频前 N 色"时才显示 spinbox
            self._spin_top_n.setVisible(True)  # 默认显示，因为默认选中的是"高频前 N 色"
            self._mode_combo.currentIndexChanged.connect(lambda idx: self._spin_top_n.setVisible(idx == 0))

    # ---------------- Event Filter ----------------
    def eventFilter(self, obj, event):
        if self._img_lbl is not None and obj is self._img_lbl:
            if event.type() == QtCore.QEvent.Type.MouseButtonPress:
                self._start(event)
            elif event.type() == QtCore.QEvent.Type.MouseMove and self._band:
                self._move(event)
            elif event.type() == QtCore.QEvent.Type.MouseButtonRelease and self._band:
                self._end(event)
        return super().eventFilter(obj, event)

    def _start(self, ev: QtGui.QMouseEvent):
        self._origin = ev.position().toPoint()
        if not self._band:
            self._band = QtWidgets.QRubberBand(QtWidgets.QRubberBand.Shape.Rectangle, self._img_lbl)
        self._band.setGeometry(QtCore.QRect(self._origin, QtCore.QSize()))
        self._band.show()

    def _move(self, ev: QtGui.QMouseEvent):
        rect = QtCore.QRect(self._origin, ev.position().toPoint()).normalized()
        self._band.setGeometry(rect)

    def _end(self, ev: QtGui.QMouseEvent):
        self._band.hide()
        rect = self._band.geometry()
        if rect.width() < 5 or rect.height() < 5:
            return  # too small
        l = int(rect.left()/self._ratio)
        t = int(rect.top()/self._ratio)
        r = int(rect.right()/self._ratio)
        b = int(rect.bottom()/self._ratio)
        region = self._pil_img.crop((l, t, r, b))
        arr = np.array(region)
        if arr.size == 0:
            return
        # 在原图上留下绿色矩形标记
        mark = QtWidgets.QFrame(self._img_lbl)
        mark.setStyleSheet("border: 2px solid green;")
        mark.setGeometry(rect)
        mark.show()
        self._confirmed_rects.append(mark)
        if self._average:
            mean_col = tuple(int(v) for v in arr.reshape(-1,3).mean(axis=0))
            self.callback(mean_col)
            total_cnt = len(getattr(self.parent(), "manual_bg_colors", []))
            QtWidgets.QMessageBox.information(self, "已取色", f"背景颜色已设为 {mean_col}\n已累计颜色 {total_cnt} 种")
        else:
            pixels = [tuple(p) for p in arr.reshape(-1,3)]
            mode_idx = getattr(self, "_mode_combo", None).currentIndex() if hasattr(self, "_mode_combo") else 0
            if mode_idx == 0:  # 高频前 N 色
                N = getattr(self, "_spin_top_n", None).value() if hasattr(self, "_spin_top_n") else 30
                cnt = Counter(pixels)
                out_cols = [col for col, _ in cnt.most_common(N)]
            else:  # 全部颜色
                out_cols = list(set(pixels))

            self.callback(out_cols)
            total_cnt = len(getattr(self.parent(), "manual_bg_colors", []))

            QtWidgets.QMessageBox.information(
                self,
                "已取色",
                f"已选取 {len(pixels)} 像素\n本次新增颜色 {len(out_cols)} 种\n已累计颜色 {total_cnt} 种",
            )


# ---------------- 新增: 保留区域框选对话框 ----------------
class KeepAreaSelectionDialog(QtWidgets.QDialog):
    """用于拖拽矩形框选需保留区域，仅返回一次矩形坐标(l,t,r,b)。"""

    def __init__(self, image_path: str, callback: Callable[[List[Tuple[int,int,int,int]]], None], parent=None):
        super().__init__(parent)
        self.setWindowTitle("拖拽框选保留区域")
        self.callback = callback
        self._origin = QtCore.QPoint()
        self._band: QtWidgets.QRubberBand | None = None
        self._marks: list[QtWidgets.QFrame] = []  # 多个绿色标记
        self._img_lbl: QtWidgets.QLabel | None = None
        self._rects: list[tuple[int,int,int,int]] = []

        # -------- 加载并缩放图像 --------
        try:
            pix = QtGui.QPixmap(image_path)
            if pix.isNull():
                pil_img = Image.open(image_path).convert("RGB")
                self._pil_img = pil_img
                qt_img = ImageQt.ImageQt(pil_img)
                pix = QtGui.QPixmap.fromImage(qt_img)
            else:
                pil_img = Image.open(image_path).convert("RGB")
                self._pil_img = pil_img
            max_len = 800
            self._ratio = 1.0
            if max(pix.width(), pix.height()) > max_len:
                self._ratio = max_len / max(pix.width(), pix.height())
                pix = pix.scaled(int(pix.width()*self._ratio), int(pix.height()*self._ratio),
                                  QtCore.Qt.AspectRatioMode.KeepAspectRatio,
                                  QtCore.Qt.TransformationMode.SmoothTransformation)
        except Exception as e:
            QtWidgets.QMessageBox.critical(self, "错误", f"无法加载图片:\n{e}")
            self.reject()
            return

        lbl = QtWidgets.QLabel()
        lbl.setPixmap(pix)
        lbl.setFixedSize(pix.size())
        lbl.setMouseTracking(True)
        lbl.installEventFilter(self)

        lay = QtWidgets.QVBoxLayout(self)
        lay.addWidget(lbl)
        self._img_lbl = lbl

        # ----- 按钮区域 -----
        btn_box = QtWidgets.QDialogButtonBox()
        self.btn_ok = btn_box.addButton("确定", QtWidgets.QDialogButtonBox.ButtonRole.AcceptRole)
        self.btn_reset = btn_box.addButton("取消", QtWidgets.QDialogButtonBox.ButtonRole.ActionRole)
        btn_box.addButton("关闭", QtWidgets.QDialogButtonBox.ButtonRole.RejectRole)
        self.btn_ok.setEnabled(False)
        self.btn_reset.setEnabled(False)
        lay.addWidget(btn_box)

        self.btn_ok.clicked.connect(self._confirm)
        self.btn_reset.clicked.connect(self._reset_selection)
        btn_box.rejected.connect(self.reject)

    # ---------------- Event Filter ----------------
    def eventFilter(self, obj, event):
        if self._img_lbl is not None and obj is self._img_lbl:
            if event.type() == QtCore.QEvent.Type.MouseButtonPress:
                self._start(event)
            elif event.type() == QtCore.QEvent.Type.MouseMove and self._band:
                self._move(event)
            elif event.type() == QtCore.QEvent.Type.MouseButtonRelease and self._band:
                self._end(event)
        return super().eventFilter(obj, event)

    def _start(self, ev: QtGui.QMouseEvent):
        self._origin = ev.position().toPoint()
        if not self._band:
            self._band = QtWidgets.QRubberBand(QtWidgets.QRubberBand.Shape.Rectangle, self._img_lbl)
        self._band.setGeometry(QtCore.QRect(self._origin, QtCore.QSize()))
        self._band.show()

    def _move(self, ev: QtGui.QMouseEvent):
        rect = QtCore.QRect(self._origin, ev.position().toPoint()).normalized()
        self._band.setGeometry(rect)

    def _end(self, ev: QtGui.QMouseEvent):
        self._band.hide()
        rect = self._band.geometry()
        if rect.width() < 5 or rect.height() < 5:
            return  # too small

        # 画持久绿色矩形
        mk = QtWidgets.QFrame(self._img_lbl)
        mk.setStyleSheet("border: 2px solid green;")
        mk.setGeometry(rect)
        mk.show()
        self._marks.append(mk)

        # 保存原图坐标
        l = int(rect.left()/self._ratio)
        t = int(rect.top()/self._ratio)
        r = int(rect.right()/self._ratio)
        b = int(rect.bottom()/self._ratio)
        self._rects.append((l, t, r, b))

        # 启用按钮
        self.btn_ok.setEnabled(True)
        self.btn_reset.setEnabled(True)

    # ---------------- 内部槽 ----------------
    def _reset_selection(self):
        for m in self._marks:
            m.deleteLater()
        self._marks.clear()
        self._rects.clear()
        self.btn_ok.setEnabled(False)
        self.btn_reset.setEnabled(False)

    def _confirm(self):
        if not self._rects:
            return
        self.callback(self._rects)
        self.accept()


# --------------------- 主窗口 ---------------------
class TransparencyTool(QtWidgets.QWidget):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("图片背景透明化工具")
        self.resize(1100, 700)
        self.root = RootAdapter()

        # ---- 状态 ----
        self.image_paths: List[str] = []
        self.batch_save_dir: str | None = None
        self.bg_color: Tuple[int, int, int] = (255,255,255)
        self.bg_image_path: str | None = None
        self.manual_bg_colors: List[Tuple[int,int,int]] = []
        self.keep_areas: List[Tuple[int,int,int,int]] = []  # 多矩形列表

        # ---- tkinter 变量兼容 ----
        self.bg_mode_var = SimpleVar("transparent")
        self.detect_mode_var = SimpleVar("auto")
        self.alg_mode_var = SimpleVar("color")
        self.auto_crop_var = SimpleVar(False)

        # ---- UI 构建 ----
        self._build_ui()

    # ==========================================================
    def _build_ui(self):
        outer = QtWidgets.QVBoxLayout(self)
        outer.setContentsMargins(10,10,10,10)
        outer.setSpacing(6)

        # ------------ 主体 Splitter ------------
        splitter = QtWidgets.QSplitter(QtCore.Qt.Orientation.Horizontal)
        outer.addWidget(splitter, 1)
        # 在添加完子控件后再设置左右宽度比例
        # ============ 左侧: 图片列表 ============
        left_widget = QtWidgets.QWidget()
        left_lay = QtWidgets.QVBoxLayout(left_widget)
        left_lay.setContentsMargins(0,0,0,0)
        left_lay.setSpacing(4)

        self.file_list = QtWidgets.QListWidget()
        self.file_list.setMinimumWidth(240)
        left_lay.addWidget(self.file_list, 1)

        self.btn_browse = QtWidgets.QPushButton("浏览图片…")
        self.btn_browse.clicked.connect(self.select_images)
        left_lay.addWidget(self.btn_browse)

        splitter.addWidget(left_widget)
        splitter.setStretchFactor(0, 1)

        # ============ 右侧: 设置区域 ============
        right_widget = QtWidgets.QWidget()
        splitter.addWidget(right_widget)
        splitter.setStretchFactor(1, 2)
        # 现在再设置默认宽度比例 (左 280px, 右 剩余)
        splitter.setSizes([280, 700])
        right = QtWidgets.QVBoxLayout(right_widget)
        right.setContentsMargins(0,0,0,0)
        right.setSpacing(8)

        # ----- 参数 sliders -----
        param_grp = QtWidgets.QGroupBox("抠图参数")
        right.addWidget(param_grp)
        pv = QtWidgets.QVBoxLayout(param_grp)

        self.tolerance_slider, self.lbl_tol = self._add_slider_row(pv, "背景去除程度", 0, 100, 9)
        self.feather_slider, self.lbl_fea = self._add_slider_row(pv, "边缘柔和度", 0, 20, 2)
        self.edge_slider, self.lbl_edge = self._add_slider_row(pv, "边缘精细度", 1, 15, 2)
        self.tolerance = SliderWrapper(self.tolerance_slider)
        self.feather = SliderWrapper(self.feather_slider)
        self.edge_clarity = SliderWrapper(self.edge_slider)

        # ----- 背景 & 算法 选项 -----
        bg_grp = QtWidgets.QGroupBox("背景自定义 / 背景识别 / 处理模式")
        right.addWidget(bg_grp)
        glay = QtWidgets.QGridLayout(bg_grp)
        # 预留列宽，保证左侧 8 个圆形单选按钮位置固定
        glay.setColumnMinimumWidth(4, 120)
        glay.setColumnMinimumWidth(5, 120)
        row = 0
        # 背景模式 (输出)
        glay.addWidget(QtWidgets.QLabel("输出背景:"), row, 0)
        self.rb_transparent = QtWidgets.QRadioButton("透明")
        self.rb_transparent.setChecked(True)
        self.rb_solid = QtWidgets.QRadioButton("颜色选取")
        self.rb_image = QtWidgets.QRadioButton("自定义图片")
        # 独立互斥组 - 输出背景
        self.grp_output = QtWidgets.QButtonGroup(self)
        for i, rb in enumerate((self.rb_transparent, self.rb_solid, self.rb_image)):
            self.grp_output.addButton(rb)
            rb.toggled.connect(lambda _, m=["transparent","solid","image"][i]: self._set_bg_mode(m))
            glay.addWidget(rb, row, i+1)

        # 纯色/图片辅助按钮与输出背景同一行，放到更右侧
        self.btn_choose_color = QtWidgets.QPushButton("选择颜色")
        self.btn_choose_color.clicked.connect(self.choose_bg_color)
        self.btn_choose_color.hide()
        glay.addWidget(self.btn_choose_color, row, 4)

        self.btn_sample_bg = QtWidgets.QPushButton("取样背景色")
        self.btn_sample_bg.clicked.connect(self.sample_bg_color)
        self.btn_sample_bg.hide()
        glay.addWidget(self.btn_sample_bg, row, 5)

        self.btn_choose_image = QtWidgets.QPushButton("选择图片")
        self.btn_choose_image.clicked.connect(self.choose_bg_image)
        self.btn_choose_image.hide()
        glay.addWidget(self.btn_choose_image, row, 4)  # 与 choose_color 位置复用

        row += 1
        # 背景识别模式
        glay.addWidget(QtWidgets.QLabel("背景识别:"), row, 0)
        self.rb_auto = QtWidgets.QRadioButton("智能自动")
        self.rb_auto.setChecked(True)
        self.rb_sample = QtWidgets.QRadioButton("吸管取样")
        # 独立互斥组 - 背景识别
        self.grp_detect = QtWidgets.QButtonGroup(self)
        for i, rb in enumerate((self.rb_auto, self.rb_sample)):
            self.grp_detect.addButton(rb)
            rb.toggled.connect(lambda _, m=["auto","sample"][i]: self._set_detect_mode(m))
            glay.addWidget(rb, row, i+1)
        self.btn_pick_color = QtWidgets.QPushButton("取样颜色")
        self.btn_pick_color.setEnabled(False)
        self.btn_pick_color.clicked.connect(self.open_color_picker)
        glay.addWidget(self.btn_pick_color, row, 4)
        self.btn_reset_sample = QtWidgets.QPushButton("重置取样色")
        self.btn_reset_sample.clicked.connect(self.reset_samples)
        glay.addWidget(self.btn_reset_sample, row, 5)

        # 默认隐藏取样按钮(仅在"吸管取样"模式显示)
        self.btn_pick_color.hide()
        self.btn_reset_sample.hide()

        row += 1
        # 处理模式 (更易懂)
        glay.addWidget(QtWidgets.QLabel("处理模式:"), row, 0)
        self.rb_color_alg = QtWidgets.QRadioButton("标准(更快)")
        self.rb_color_alg.setChecked(True)
        self.rb_flood_alg = QtWidgets.QRadioButton("精细(更慢)")
        self.rb_ai_alg = QtWidgets.QRadioButton("AI抠图(较慢)")
        # 互斥组 - 算法
        self.grp_alg = QtWidgets.QButtonGroup(self)
        alg_buttons = (
            (self.rb_color_alg, "color"),
            (self.rb_flood_alg, "flood"),
            (self.rb_ai_alg, "ai"),
        )
        for i, (rb, mode) in enumerate(alg_buttons):
            self.grp_alg.addButton(rb)
            rb.toggled.connect(lambda _, m=mode: self._set_alg_mode(m))
            glay.addWidget(rb, row, i+1)

        # ----- 自动裁剪选项 -----
        self.chk_crop = QtWidgets.QCheckBox("处理后裁剪透明边")
        self.chk_crop.stateChanged.connect(lambda s: self.auto_crop_var.set(bool(s)))
        right.addWidget(self.chk_crop)

        # ----- 保留区域框选按钮 (仅单张图片可用) -----
        self.btn_keep_area = QtWidgets.QPushButton("框选保留区域")
        self.btn_keep_area.setEnabled(False)
        self.btn_keep_area.clicked.connect(self.select_keep_area)
        right.addWidget(self.btn_keep_area)

        # ----- 开始处理按钮 (居中 & 更大) -----
        self.btn_process = QtWidgets.QPushButton("开始处理")
        self.btn_process.setEnabled(False)
        self.btn_process.setFixedSize(140, 40)  # 增大尺寸
        font = self.btn_process.font()
        font.setPointSize(font.pointSize() + 2)
        self.btn_process.setFont(font)
        self.btn_process.clicked.connect(self.process_images)
        right.addWidget(self.btn_process, alignment=QtCore.Qt.AlignmentFlag.AlignHCenter)

        # ----- 进度条 -----
        self.progress_label = QtWidgets.QLabel()
        outer.addWidget(self.progress_label)
        self.progress_bar = QtWidgets.QProgressBar()
        self.progress_bar.setTextVisible(False)
        self.progress_bar.hide()
        outer.addWidget(self.progress_bar)

    # ---------- slider 工具 ----------
    def _add_slider_row(self, parent_layout: QtWidgets.QVBoxLayout, text:str, mn:int,mx:int,init:int):
        lay = QtWidgets.QHBoxLayout()
        parent_layout.addLayout(lay)
        lab = QtWidgets.QLabel(text+":")
        lay.addWidget(lab)
        slider = QtWidgets.QSlider(QtCore.Qt.Orientation.Horizontal)
        slider.setRange(mn,mx)
        slider.setValue(init)
        lay.addWidget(slider,1)
        val_lab = QtWidgets.QLabel(str(init))
        lay.addWidget(val_lab)
        slider.valueChanged.connect(lambda v, lbl=val_lab: lbl.setText(str(v)))
        return slider, val_lab

    # ==========================================================
    # ---------------- 状态更新 ----------------
    def _set_bg_mode(self, mode:str):
        self.bg_mode_var.set(mode)
        self.btn_choose_color.setVisible(mode=="solid")
        self.btn_sample_bg.setVisible(mode=="solid")
        self.btn_choose_image.setVisible(mode=="image")

    def _set_detect_mode(self, mode:str):
        self.detect_mode_var.set(mode)
        show_sample_btns = (mode == "sample")
        # 仅在取样模式下显示/启用按钮
        self.btn_pick_color.setVisible(show_sample_btns)
        self.btn_reset_sample.setVisible(show_sample_btns)
        self.btn_pick_color.setEnabled(show_sample_btns and bool(self.image_paths))
        # -------- AI 算法可用性 --------
        ai_enabled = (mode == "auto")
        self.rb_ai_alg.setEnabled(ai_enabled)
        if not ai_enabled and self.rb_ai_alg.isChecked():
            # 若当前选中了 AI，但用户切到吸管，则自动切回速度较快的标准算法
            self.rb_color_alg.setChecked(True)

    def _set_alg_mode(self, mode:str):
        self.alg_mode_var.set(mode)

    # ==========================================================
    # ---------------- 文件选择 ----------------
    def select_images(self):
        files, _ = QtWidgets.QFileDialog.getOpenFileNames(self, "选择图片", "",
            "图片文件 (*.png *.jpg *.jpeg *.bmp *.gif *.tiff *.webp)")
        if not files:
            return
        self.image_paths = files
        self.file_list.clear()
        for p in files:
            disp = os.path.basename(p)
            self.file_list.addItem(disp)
        self.btn_process.setEnabled(True)
        self.btn_pick_color.setEnabled(self.detect_mode_var.get()=="sample")
        self.rb_ai_alg.setEnabled(self.detect_mode_var.get()=="auto")
        # ---- 仅单张图片时启用框选按钮 ----
        self.btn_keep_area.setEnabled(len(files)==1)
        self.keep_areas.clear()  # 切换图片时重置

    # ==========================================================
    # ---------------- 进度条 ----------------
    def _show_progress(self, msg:str="处理中..."):
        self.progress_label.setText(msg)
        self.progress_bar.show()
        self.progress_bar.setRange(0,0)
        self.btn_process.setEnabled(False)
        QtWidgets.QApplication.processEvents()

    def _hide_progress(self):
        self.progress_bar.hide()
        self.progress_label.clear()
        self.btn_process.setEnabled(True)
        QtWidgets.QApplication.processEvents()

    # ==========================================================
    # ---------------- 图片处理 ----------------
    def process_images(self):
        if not self.image_paths:
            return
        # 批量输出目录
        if len(self.image_paths) > 1:
            out_dir = QtWidgets.QFileDialog.getExistingDirectory(self, "选择输出文件夹")
            if not out_dir:
                return
            self.batch_save_dir = out_dir
        else:
            self.batch_save_dir = None
        threading.Thread(target=self._process_thread, daemon=True).start()

    def _process_thread(self):
        total = len(self.image_paths)
        for idx, path in enumerate(self.image_paths,1):
            self._process_single(path, idx, total)
        self.root.after(0, self._hide_progress)

    def _process_single(self, img_path:str, idx:int,total:int):
        prefix = f"({idx}/{total}) " if total>1 else ""
        self.root.after(0, lambda: self._show_progress(prefix+"读取图片..."))
        try:
            with open(img_path, 'rb') as f:
                buf = np.frombuffer(f.read(), np.uint8)
            img_bgr = cv2.imdecode(buf, cv2.IMREAD_COLOR)
        except Exception as e:
            self.root.after(0, lambda: self._error(str(e)))
            return
        if img_bgr is None:
            self.root.after(0, lambda: self._error("无法读取图片"))
            return
        img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
        h,w = img_rgb.shape[:2]

        # --------- 内存优化: 对超高分辨率图片做临时缩放 ---------
        MAX_PIXELS = 1_500_000  # 约 150 万像素，可根据需要调整
        scale_ratio = 1.0
        if h * w > MAX_PIXELS:
            scale_ratio = (MAX_PIXELS / (h * w)) ** 0.5

        self.root.after(0, lambda: self._show_progress(prefix+"分析背景颜色..."))
        edge = []
        edge.extend(img_rgb[0,:,:].reshape(-1,3))
        edge.extend(img_rgb[h-1,:,:].reshape(-1,3))
        edge.extend(img_rgb[:,0,:].reshape(-1,3))
        edge.extend(img_rgb[:,w-1,:].reshape(-1,3))
        edge = [tuple(p) for p in edge]
        cnt = Counter(edge)
        bg_colors = [c for c,n in cnt.items() if n/len(edge)*100>=1] or [cnt.most_common(1)[0][0]]
        if self.detect_mode_var.get()=="sample":
            if not self.manual_bg_colors:
                self.root.after(0, lambda: self._error("请先取样背景颜色"))
                return
            bg_colors = self.manual_bg_colors

        tol = self.tolerance.get()
        fea = self.feather.get()
        edg = self.edge_clarity.get()

        self.root.after(0, lambda: self._show_progress(prefix+"抠图处理中..."))
        alg_mode = self.alg_mode_var.get()

        # 按需缩放输入，减少算法阶段内存压力（AI 模式除外）
        img_for_alg = img_rgb
        if alg_mode != "ai" and scale_ratio != 1.0:
            new_w = max(1, int(w * scale_ratio))
            new_h = max(1, int(h * scale_ratio))
            img_for_alg = cv2.resize(img_rgb, (new_w, new_h), interpolation=cv2.INTER_AREA)

        # ===== 保留区域优先处理 =====
        if self.keep_areas:
            alpha = np.zeros((h, w), dtype=np.uint8)

            if alg_mode == "ai":
                # 逐矩形直接调用 AI 算法（原尺寸）
                for l, t, r, b in self.keep_areas:
                    l = max(0, min(l, w-1)); r = max(0, min(r, w-1))
                    t = max(0, min(t, h-1)); b = max(0, min(b, h-1))
                    sub_img = img_rgb[t:b+1, l:r+1]
                    try:
                        sub_alpha = u2net_alpha(sub_img)
                    except Exception as e:
                        self.root.after(0, lambda: self._error(str(e)))
                        return
                    alpha[t:b+1, l:r+1] = sub_alpha
            else:
                for l, t, r, b in self.keep_areas:
                    l = max(0, min(l, w-1)); r = max(0, min(r, w-1))
                    t = max(0, min(t, h-1)); b = max(0, min(b, h-1))

                    if scale_ratio != 1.0:
                        ls, rs = int(l*scale_ratio), int(r*scale_ratio)
                        ts, bs = int(t*scale_ratio), int(b*scale_ratio)
                        sub_img_alg = img_for_alg[ts:bs+1, ls:rs+1]
                    else:
                        sub_img_alg = img_for_alg[t:b+1, l:r+1]

                    if sub_img_alg.size == 0:
                        continue

                    if alg_mode == "flood":
                        seed_colors = [] if self.detect_mode_var.get()=="auto" else bg_colors
                        sub_alpha_small = floodfill_alpha(cv2.cvtColor(sub_img_alg, cv2.COLOR_RGB2RGBA), seed_colors, tol, fea, edg)
                    else:  # 标准(color)模式
                        sub_alpha_small = distance_alpha(sub_img_alg, bg_colors, tol, fea, edg)

                    if scale_ratio != 1.0:
                        sub_alpha = cv2.resize(sub_alpha_small, (r-l+1, b-t+1), interpolation=cv2.INTER_LINEAR)
                    else:
                        sub_alpha = sub_alpha_small

                    alpha[t:b+1, l:r+1] = sub_alpha

            # 生成 RGBA 结果，用于后续流程
            res = np.zeros((h, w, 4), dtype=np.uint8)
            res[:, :, :3] = img_rgb
            res[:, :, 3] = alpha

        else:
            # ====== 原整图处理逻辑 ======
            if alg_mode == "flood":
                # 当背景识别为"自动"时，传递空列表让算法自行检测；
                # 取样模式下则传入用户取样的颜色作为 flood-fill 种子
                seed_colors = [] if self.detect_mode_var.get()=="auto" else bg_colors
                alpha_small = floodfill_alpha(cv2.cvtColor(img_for_alg, cv2.COLOR_RGB2RGBA), seed_colors, tol, fea, edg)
            elif alg_mode == "ai":
                try:
                    alpha_small = u2net_alpha(img_rgb)  # AI 模式直接使用原尺寸
                except Exception as e:
                    self.root.after(0, lambda: self._error(str(e)))
                    return
            else:
                alpha_small = distance_alpha(img_for_alg, bg_colors, tol, fea, edg)

            # 若做了缩放，则把掩码还原到原尺寸
            if scale_ratio != 1.0 and alg_mode != "ai":
                alpha = cv2.resize(alpha_small, (w, h), interpolation=cv2.INTER_LINEAR)
            else:
                alpha = alpha_small

            res = np.zeros((h, w, 4), dtype=np.uint8)
            res[:, :, :3] = img_rgb
            res[:, :, 3] = alpha

        # ---- 若有保留区域限制，矩形外全部设为透明 ----
        if self.keep_areas:
            mask = np.zeros_like(alpha, dtype=np.uint8)
            for l,t,r,b in self.keep_areas:
                l = max(0, min(l, w-1))
                r = max(0, min(r, w-1))
                t = max(0, min(t, h-1))
                b = max(0, min(b, h-1))
                mask[t:b+1, l:r+1] = 255
            res[:,:,3] = cv2.bitwise_and(res[:,:,3], mask)
            alpha = res[:,:,3]

        if self.auto_crop_var.get():
            coords = np.column_stack(np.where(alpha>0))
            if coords.size:
                ymin,xmin = coords.min(axis=0)
                ymax,xmax = coords.max(axis=0)
                res = res[ymin:ymax+1, xmin:xmax+1, :]

        mode = self.bg_mode_var.get()
        if mode=="transparent":
            out_img = Image.fromarray(res)
        else:
            fg = Image.fromarray(res)
            if mode=="solid":
                bg = Image.new('RGBA', fg.size, self.bg_color+(255,))
            else:
                if not self.bg_image_path or not os.path.exists(self.bg_image_path):
                    self.root.after(0, lambda: self._error("请先选择背景图片"))
                    return
                bg = Image.open(self.bg_image_path).convert('RGBA').resize(fg.size, Image.LANCZOS)
            out_img = Image.alpha_composite(bg, fg).convert('RGB')

        if self.batch_save_dir:
            prefix_name = "transparent_" if mode=="transparent" else "output_"
            save_path = os.path.join(self.batch_save_dir, prefix_name+os.path.basename(img_path))
            try:
                out_img.save(save_path, 'PNG')
            except Exception as e:
                self.root.after(0, lambda: self._error(f"保存失败:{e}"))
        else:
            self.root.after(0, lambda img=out_img, m=mode: self._save_single(img, m))

    # ---------------- save 单张 ----------------
    def _save_single(self, img:Image.Image, mode:str):
        self._hide_progress()
        default = "transparent.png" if mode=="transparent" else "output.png"
        path, _ = QtWidgets.QFileDialog.getSaveFileName(self, "保存文件", default, "PNG files (*.png)")
        if not path:
            return
        self._show_progress("保存图片...")
        def _save():
            try:
                img.save(path, 'PNG')
                self.root.after(0, self._hide_progress)
            except Exception as e:
                self.root.after(0, lambda: self._error(f"保存失败:{e}"))
        threading.Thread(target=_save, daemon=True).start()

    # ---------------- 错误 ----------------
    def _error(self, msg:str):
        self._hide_progress()
        QtWidgets.QMessageBox.critical(self, "错误", msg)

    # ==========================================================
    # ---------------- 取样 ----------------
    def open_color_picker(self):
        if self.detect_mode_var.get()!="sample" or not self.image_paths:
            return
        def cb(cols:List[Tuple[int,int,int]]):
            for c in cols:
                if c not in self.manual_bg_colors:
                    self.manual_bg_colors.append(c)
        RegionSelectionDialog(self.image_paths[0], cb, average=False, parent=self).exec()

    def reset_samples(self):
        self.manual_bg_colors.clear()
        QtWidgets.QMessageBox.information(self, "已重置", "已清空手动取样颜色")

    # ---- 背景颜色/图片选择 ----
    def choose_bg_color(self):
        col = QtWidgets.QColorDialog.getColor(QtGui.QColor(*self.bg_color), self, "选择颜色")
        if col.isValid():
            self.bg_color = (col.red(), col.green(), col.blue())

    def sample_bg_color(self):
        path, _ = QtWidgets.QFileDialog.getOpenFileName(self, "选择参考图片", "", "图片文件 (*.png *.jpg *.jpeg *.bmp *.gif *.tiff *.webp)")
        if not path:
            return
        def cb(mean_col:Tuple[int,int,int]):
            self.bg_color = mean_col
        RegionSelectionDialog(path, cb, average=True, parent=self).exec()

    def choose_bg_image(self):
        path, _ = QtWidgets.QFileDialog.getOpenFileName(self, "选择背景图片", "", "图片文件 (*.png *.jpg *.jpeg *.bmp *.gif *.tiff *.webp)")
        if path:
            self.bg_image_path = path

    # ==========================================================
    # ---------------- 保留区域选择 ----------------
    def select_keep_area(self):
        if not self.image_paths or len(self.image_paths)!=1:
            return
        def cb(rects:List[Tuple[int,int,int,int]]):
            self.keep_areas = rects
            QtWidgets.QMessageBox.information(self, "已设置", f"已设定 {len(rects)} 个保留区域")
        KeepAreaSelectionDialog(self.image_paths[0], cb, parent=self).exec()


# ---------------- 入口 ----------------
if __name__ == "__main__":
    app = QtWidgets.QApplication(sys.argv)
    QtWidgets.QApplication.setStyle("Fusion")
    win = TransparencyTool()
    win.show()
    sys.exit(app.exec())
